#include "common.h"
#include "auxiliary.h"

Matcp Dx(Matcp u);
Matcp Dxt(Matcp u);
Matcp Dy(Matcp u);
Matcp Dyt(Matcp u);
void shrink2(Matcp* xout, Matcp* yout, Matcp xin, Matcp yin, rltype lambda);

Matcp split_bregman(Matcp R, Matcp f, rltype mu, rltype lambda, rltype gamma, int nInner, int nOuter)
{
	const int frows = f.rows;
	const int fcols = f.cols;
	const rltype scale = sqrt(rltype(frows*fcols));
	const rltype scalercp = 1.0/scale;
	const Matcp f0 = f.clone();

	Matcp u(frows, fcols, 0.0); // function output
	Matcp x(frows, fcols, 0.0);
	Matcp y(frows, fcols, 0.0);
	Matcp bx(frows, fcols, 0.0);
	Matcp by(frows, fcols, 0.0);
	Matcp uker(frows, fcols, 0.0);
	Matcp murf;

	// build kernels
	murf = scalercp * mu * elmult(R, f);
	idft(murf, murf);


	uker(0,0) = 4.0;
	uker(1,0) = -1.0;
	uker(0,1) = -1.0;
	uker(frows-1, 0) = -1.0;
	uker(0, fcols-1) = -1.0;
	dft(uker, uker);
	uker = addscalar(mu * R + lambda * uker, gamma);


	// do the reconstruction
	Matcp rhs, uf, dx, dy;
	for(int outer = 0; outer < nOuter; outer++)
	{
		for(int inner = 0; inner < nInner; inner++)
		{
			// update u
			rhs = murf + lambda * Dxt(x-bx) + lambda * Dyt(y-by) + gamma * u;
			dft(rhs, rhs);
			rhs = (scalercp * scalercp) * eldiv(rhs, uker);
			idft(rhs, u);

			// update x and y
			dx = Dx(u);
			dy = Dy(u);
			shrink2(&x, &y, dx+bx, dy+by, 1.0/lambda);

			// update bregman parameters
			bx += dx - x;
			by += dy - y;
		}
		dft(u, uf);
		f += f0 - scalercp * elmult(R, uf);
		idft(mu * scalercp * elmult(R, f), murf);
	}
	return u;
}

Matcp Dx(Matcp u)
// Forward difference in the x-direction.
{
	const int urows = u.rows;
	const int ucols = u.cols;
	
	Matcp ux = Matcp(urows, ucols, 0.0);
	ux.col(0) = u.col(0) - u.col(ucols-1);
	ux(Range::all(), Range(1, ucols)) = u(Range::all(), Range(1, ucols)) - u(Range::all(), Range(0, ucols-1));

	return ux;
}

Matcp Dxt(Matcp u)
// Backward difference in the x-direction.
{
	const int urows = u.rows;
	const int ucols = u.cols;

	Matcp uxt = Matcp(urows, ucols, 0.0);
	uxt(Range::all(), Range(0, ucols-1)) = u(Range::all(), Range(0, ucols-1)) - u(Range::all(), Range(1, ucols));
	uxt.col(ucols-1) = u.col(ucols-1) - u.col(0);

	return uxt;
}

Matcp Dy(Matcp u)
// Forward difference in the y-direction.
{
	const int urows = u.rows;
	const int ucols = u.cols;

	Matcp uy = Matcp(urows, ucols, 0.0);
	uy.row(0) = u.row(0) - u.row(urows-1);
	uy(Range(1, urows), Range::all()) = u(Range(1, urows), Range::all()) - u(Range(0, urows-1), Range::all());

	return uy;
}

Matcp Dyt(Matcp u)
// Backward difference in the y-direction
{
	const int urows = u.rows;
	const int ucols = u.cols;

	Matcp uyt = Matcp(urows, ucols, 0.0);
	uyt(Range(0, urows-1), Range::all()) = u(Range(0, urows-1), Range::all()) - u(Range(1, urows), Range::all());
	uyt.row(urows-1) = u.row(urows-1) - u.row(0);

	return uyt;
}

void shrink2(Matcp* xout, Matcp* yout, Matcp x, Matcp y, rltype lambda)
// This should have a better explanation
{
	Matrl s;
	Matrl xabs = cvabs(x);
	pow(xabs, 2, xabs);

	Matrl yabs = cvabs(y);
	pow(yabs, 2, yabs);
	pow(xabs + yabs, 0.5, s);

	Matrl ss = s - lambda;
	ss = thresholdtz64(ss, 0);
	s = s + thresholdbi64(s, lambda, 1);

	ss = eldiv(ss, s);


	Matcp sscp = merge2(ss, Mat::zeros(ss.size(), CV_RL));
	*xout = elmult(x, sscp);
	*yout = elmult(y, sscp);
}